package com.lost.octopus.es.processor.config.es.alias;

import lombok.extern.log4j.Log4j2;
import org.elasticsearch.index.query.IdsQueryBuilder;
import org.springframework.data.elasticsearch.core.AliasIndexOperations;
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
import org.springframework.data.elasticsearch.repository.support.ElasticsearchEntityInformation;
import org.springframework.data.elasticsearch.repository.support.SimpleElasticsearchRepository;
import org.springframework.scheduling.annotation.Async;
import org.springframework.util.Assert;

import java.util.function.Function;

import static org.elasticsearch.index.query.QueryBuilders.idsQuery;

/**
 * 支持别名的es仓储类
 * @author zhangbh
 * @date 2022/10/31 0031 11:34
 */
@Log4j2
public class SimpleAliasElasticsearchRepository<T, ID> extends SimpleElasticsearchRepository<T, ID> implements AliasElasticsearchRepository<T, ID> {

    public SimpleAliasElasticsearchRepository(ElasticsearchEntityInformation<T, ID> metadata, ElasticsearchOperations operations) {
        super(metadata, operations);
    }


    @Override
    public <S extends T> Iterable<S> saveAll(Iterable<S> entities, IndexCoordinates indexCoordinates) {
        Assert.notNull(entities, "Cannot insert 'null' as a List.");
        if (indexCoordinates == null) {
            return saveAll(entities);
        }

        // 保存
        operations.save(entities, indexCoordinates);
        // 刷新指定索引
        validAliasOperation(indexOperations -> {
            indexOperations.refresh(indexCoordinates);
            return true;
        });
        return entities;
    }

    /**
     * {@link #saveAll(Iterable, IndexCoordinates)}
     * @return
     */
    @Override
    public IndexCoordinates getIndexCoordinates() {
        return operations.getIndexCoordinatesFor(entityClass);
    }

    @Override
    public void deleteAllById(Iterable<? extends ID> ids, IndexCoordinates indexCoordinates) {
        Assert.notNull(ids, "Cannot delete 'null' as a List.");
        if (indexCoordinates == null) {
            deleteAllById(ids);
            return;
        }
        IdsQueryBuilder idsQueryBuilder = idsQuery();
        for (ID id : ids) {
            idsQueryBuilder.addIds(stringIdRepresentation(id));
        }
        if (idsQueryBuilder.ids().isEmpty()) {
            return;
        }
        executeAndRefresh((OperationsCallback<Void>) operations -> {
            operations.delete(new NativeSearchQueryBuilder().withQuery(idsQueryBuilder).build(), entityClass,
                indexCoordinates);
            return null;
        });

    }

    @Override
    public IndexCoordinates getTimestampAlias(String name) {
        return validAliasOperation(indexOperations -> indexOperations.getTimestampAlias(name));
    }

    @Override
    public void addByOriginalAlias(IndexCoordinates coordinates) {
        validAliasOperation(indexOperations -> {
            indexOperations.addByOriginalAlias(coordinates);
            return true;
        });
    }

    @Override
    public boolean createWithMapping(IndexCoordinates coordinates) {
        return validAliasOperation(indexOperations -> indexOperations.createWithMapping(coordinates));
    }

    @Override
    public EsAliasRemoveInfo removeAlias() {
        return validAliasOperation(AliasIndexOperations::removeAlias);
    }

    @Override
    public boolean indexExists() {
        return validAliasOperation(indexOperations -> indexOperations.exists());
    }

    /**
     * 校验处理操作是否支持
     * @param optional 操作函数
     * @param <R>      返回类型
     * @return 返回值，若有
     */
    private <R> R validAliasOperation(Function<AliasIndexOperations, R> optional) {
        if (!(indexOperations instanceof AliasIndexOperations)) {
            throw new AliasRepositoryNoSupportedException("当前仓储初始化异常，请检查是否实现 AliasElasticsearchRepository");
        }
        return optional.apply((AliasIndexOperations) indexOperations);
    }

    @Override
    public boolean deleteIndex() {
        return indexOperations.delete();
    }

    @Override
    @Async
    public void deleteAsyncIndex(String indexName) {
        validAliasOperation(indexOperations -> indexOperations.delete(indexName));
    }
}
